from PySide6.QtWidgets import (QApplication, QWidget, QVBoxLayout, QHBoxLayout, 
                             QPushButton, QLabel, QLineEdit, QTextEdit, 
                             QProgressBar, QFileDialog, QTreeWidget, QTreeWidgetItem)
from PySide6.QtGui import QFont, QIcon
from PySide6.QtCore import Qt, QThread, Signal
import sys
import os

# 多线程扫描图片文件
class ScanThread(QThread):
    # 定义信号：传递扫描到的图片路径列表（包含完整路径和相对路径信息）
    scan_finished = Signal(list)
    # 定义信号：传递单个扫描进度
    scan_progress = Signal(str)

    def __init__(self, folder_path):
        super().__init__()
        self.folder_path = folder_path
        self.image_extensions = ('.jpg', '.jpeg', '.png', '.webp')

    def run(self):
        image_info = []
        # 递归扫描文件夹
        for root, _, files in os.walk(self.folder_path):
            for file in files:
                if file.lower().endswith(self.image_extensions):
                    full_path = os.path.join(root, file)
                    # 计算相对路径（用于树状显示层级）
                    relative_path = os.path.relpath(full_path, self.folder_path)
                    image_info.append((full_path, relative_path))
                    self.scan_progress.emit(f"发现图片：{relative_path}")
        self.scan_finished.emit(image_info)

class MainWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.image_list = []  # 存储扫描到的图片完整路径
        self.init_ui()

    def init_ui(self):
        self.setWindowTitle("图片处理工具")
        self.setFixedSize(900, 650)  # 加宽窗口以适应树状显示

        # 主布局
        main_layout = QVBoxLayout()
        main_layout.setContentsMargins(20, 20, 20, 20)
        main_layout.setSpacing(10)

        # 字体设置
        font = QFont()
        font.setPointSize(10)

        # 1. 选择图片文件夹区域
        folder_layout = QHBoxLayout()
        folder_label = QLabel("图片文件夹:")
        folder_label.setFont(font)
        self.folder_entry = QLineEdit()
        self.folder_entry.setFont(font)
        folder_btn = QPushButton("选择")
        folder_btn.setFont(font)
        folder_btn.clicked.connect(self.select_image_folder)
        folder_layout.addWidget(folder_label)
        folder_layout.addWidget(self.folder_entry)
        folder_layout.addWidget(folder_btn)

        # 2. 选择贴图文件区域
        sticker_layout = QHBoxLayout()
        sticker_label = QLabel("贴图文件:")
        sticker_label.setFont(font)
        self.sticker_entry = QLineEdit()
        self.sticker_entry.setFont(font)
        sticker_btn = QPushButton("选择")
        sticker_btn.setFont(font)
        sticker_btn.clicked.connect(self.select_sticker_path)
        sticker_layout.addWidget(sticker_label)
        sticker_layout.addWidget(self.sticker_entry)
        sticker_layout.addWidget(sticker_btn)

        # 3. 图片树状显示区域（新增）
        tree_label = QLabel("待处理图片（树状结构）:")
        tree_label.setFont(font)
        self.image_tree = QTreeWidget()
        self.image_tree.setFont(font)
        self.image_tree.setHeaderLabel("文件结构")  # 设置树状控件标题
        self.image_tree.setAlternatingRowColors(True)  # 行交替颜色，增强可读性

        # 4. 开始处理按钮
        start_btn = QPushButton("开始处理")
        start_btn.setFont(font)
        start_btn.clicked.connect(self.start_processing)

        # 5. 日志信息框
        log_label = QLabel("处理日志:")
        log_label.setFont(font)
        self.info_text = QTextEdit()
        self.info_text.setReadOnly(True)
        self.info_text.setFont(font)

        # 6. 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setAlignment(Qt.AlignHCenter)

        # 添加所有组件到主布局
        main_layout.addLayout(folder_layout)
        main_layout.addLayout(sticker_layout)
        main_layout.addWidget(tree_label)
        main_layout.addWidget(self.image_tree)
        main_layout.addWidget(start_btn)
        main_layout.addWidget(log_label)
        main_layout.addWidget(self.info_text)
        main_layout.addWidget(self.progress_bar)

        self.setLayout(main_layout)

    def select_image_folder(self):
        folder_path = QFileDialog.getExistingDirectory(self, "选择图片文件夹")
        if folder_path:
            self.folder_entry.setText(folder_path)
            self.info_text.append(f"开始扫描文件夹: {folder_path}")
            # 清空之前的树状结构
            self.image_tree.clear()
            self.image_list = []
            # 启动多线程扫描
            self.scan_thread = ScanThread(folder_path)
            self.scan_thread.scan_finished.connect(self.update_image_tree)
            self.scan_thread.scan_progress.connect(self.update_scan_log)
            self.scan_thread.start()

    def select_sticker_path(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择贴图文件", "", "Image files (*.jpg *.png *.jpeg *.webp)"
        )
        if file_path:
            self.sticker_entry.setText(file_path)
            self.info_text.append(f"已选择贴图文件: {os.path.basename(file_path)}")

    def update_scan_log(self, msg):
        self.info_text.append(msg)

    def update_image_tree(self, image_info):
        # 扫描完成后，更新树状结构
        root_path = self.folder_entry.text()
        root_item = QTreeWidgetItem([os.path.basename(root_path)])
        self.image_tree.addTopLevelItem(root_item)
        
        # 存储所有节点，避免重复创建
        node_cache = {root_path: root_item}
        
        for full_path, relative_path in image_info:
            self.image_list.append(full_path)
            current_path = root_path
            current_node = root_item
            components = relative_path.split(os.sep)
            
            # 逐级创建文件夹节点
            for i, component in enumerate(components[:-1]):  # 排除文件名
                current_path = os.path.join(current_path, component)
                if current_path not in node_cache:
                    new_node = QTreeWidgetItem([component])
                    current_node.addChild(new_node)
                    node_cache[current_path] = new_node
                    current_node = new_node
                else:
                    current_node = node_cache[current_path]
            
            # 添加文件节点
            file_name = components[-1]
            file_node = QTreeWidgetItem([file_name])
            current_node.addChild(file_node)
        
        # 自动展开顶层节点
        root_item.setExpanded(True)
        self.info_text.append(f"扫描完成，共发现 {len(image_info)} 张图片")

    def start_processing(self):
        if not self.image_list:
            self.info_text.append("请先选择图片文件夹并等待扫描完成")
            return
        if not self.sticker_entry.text():
            self.info_text.append("请先选择贴图文件")
            return
        self.info_text.append("开始处理图片...")

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())
